package com.ahdms.es.util;


import com.ahdms.es.gm.algorithm.SM2;
import com.ahdms.es.gm.algorithm.SM2KeyPair;
import com.ahdms.es.gm.algorithm.SM3Digest;
import com.ahdms.es.gm.asn1.cipher.SM2PrivateKey;
import com.ahdms.es.gm.asn1.cipher.SM2PublicKey;
import com.ahdms.es.gm.asn1.cipher.SM2Signature;
import org.bouncycastle.asn1.ASN1Integer;
import org.bouncycastle.asn1.DERBitString;
import org.bouncycastle.asn1.x509.Certificate;
import org.bouncycastle.crypto.AsymmetricCipherKeyPair;
import org.bouncycastle.crypto.params.ECPrivateKeyParameters;
import org.bouncycastle.crypto.params.ECPublicKeyParameters;
import org.bouncycastle.math.ec.ECPoint;
import org.bouncycastle.util.Pack;

import java.math.BigInteger;

public class SM2Utils {

    private static final String DEFAULT_SIGN_ID = "1234567812345678";

    public static ECPoint validSM2ECCurvePoint(BigInteger X, BigInteger Y) {
        SM2 sm2 = SM2.getInstance();
        try {
            return sm2.getEccCurve().createPoint(X, Y);
        } catch (Exception e) {
            return null;
        }
    }

    public static boolean validSM2BigInteger(BigInteger k) {
        SM2 sm2 = SM2.getInstance();
        if (BigInteger.ONE.compareTo(k) > 0 || k.compareTo(sm2.getEccN().subtract(BigInteger.ONE)) > 0) {
            return false;
        } else {
            return true;
        }
    }

    public static ECPoint bytesToPoint(byte[] bytes) {
        SM2 sm2 = SM2.getInstance();
        byte[] formatedPubKey;
        if (bytes.length == 64) {
            //添加一字节标识，用于ECPoint解析
            formatedPubKey = new byte[65];
            formatedPubKey[0] = 0x04;
            System.arraycopy(bytes, 0, formatedPubKey, 1, bytes.length);
        } else {
            formatedPubKey = bytes;
        }

        return sm2.getEccCurve().decodePoint(formatedPubKey);
    }

    public static SM2KeyPair genKeyPair() {
        SM2 sm2 = SM2.getInstance();
        AsymmetricCipherKeyPair key1 = sm2.getEccKeyPairGenerator().generateKeyPair();
        ECPrivateKeyParameters ecpriv1 = (ECPrivateKeyParameters) key1.getPrivate();
        ECPublicKeyParameters ecpub1 = (ECPublicKeyParameters) key1.getPublic();
        BigInteger k1 = ecpriv1.getD();
        ECPoint pk1 = ecpub1.getQ();
        byte[] pkXBytes1 = pk1.normalize().getXCoord().getEncoded();
        byte[] pkYBytes1 = pk1.normalize().getYCoord().getEncoded();

        byte[] pubKey = new byte[64];
        System.arraycopy(pkXBytes1, 0, pubKey, 0, pkXBytes1.length);
        System.arraycopy(pkYBytes1, 0, pubKey, pkXBytes1.length, pkYBytes1.length);

        return new SM2KeyPair(k1.toByteArray(), pubKey);
    }

    public static byte[] Z(byte[] id, SM2PublicKey publicKey) {
        SM3Digest sm3 = new SM3Digest();

        SM2 sm2 = SM2.getInstance();

        int len = id.length * 8;
        sm3.update((byte) (len >> 8 & 0xFF));
        sm3.update((byte) (len & 0xFF));
        sm3.update(id, 0, id.length);

        byte[] p = ConvertorUtils.byteConvert32Bytes(sm2.getEccA());
        sm3.update(p, 0, p.length);

        p = ConvertorUtils.byteConvert32Bytes(sm2.getEccB());
        sm3.update(p, 0, p.length);

        p = ConvertorUtils.byteConvert32Bytes(sm2.getEccGx());
        sm3.update(p, 0, p.length);

        p = ConvertorUtils.byteConvert32Bytes(sm2.getEccGy());
        sm3.update(p, 0, p.length);

        p = ConvertorUtils.byteConvert32Bytes(publicKey.getXCoordinate().getPositiveValue());
        sm3.update(p, 0, p.length);

        p = ConvertorUtils.byteConvert32Bytes(publicKey.getYCoordinate().getPositiveValue());
        sm3.update(p, 0, p.length);

        byte[] md = new byte[sm3.getDigestSize()];
        sm3.doFinal(md, 0);
        return md;
    }

    public static byte[] H(byte[] z, byte[] m) {
        SM3Digest sm3 = new SM3Digest();

        sm3.update(z, 0, z.length);
        sm3.update(m, 0, m.length);
        byte[] md = new byte[32];
        sm3.doFinal(md, 0);
        return md;
    }

    public static byte[] KDF(byte[] z, int length) {
        int sm3_v = 256;
        int hashNum = (int) Math.ceil((double) length / sm3_v);
        boolean isInt = length % sm3_v == 0 ? true : false;

        SM3Digest sm3 = new SM3Digest();
        int ct = 0x1;

        byte[] tempK = new byte[hashNum * 32];

        int realByteNum = 0;
        for (int i = 1; i <= hashNum; i++) {

            byte[] ctBytes = Pack.intToBigEndian(ct);

            sm3.reset();
            sm3.update(z, 0, z.length);
            sm3.update(ctBytes, 0, ctBytes.length);

            byte[] Ha = new byte[32];
            sm3.doFinal(Ha, 0);

            ct++;

            if (hashNum == i && !isInt) { //最后一个hash，取最左边lastHashBitNum位
                int floor = (int) Math.floor((double) length / sm3_v);
                int lastHashBitNum = length - sm3_v * floor;
                int maskedByteLen = lastHashBitNum / 8;

                System.arraycopy(Ha, 0, tempK, (i - 1) * 32, maskedByteLen);
                realByteNum += maskedByteLen;

                if (0 != lastHashBitNum % 8) {
                    int maskedBitLen = 8 - lastHashBitNum % 8;
                    byte maskByte = (byte) (~((1 << (8 - maskedBitLen)) - 1));
                    tempK[(i - 1) * 32 + maskedByteLen] = (byte) (maskByte & Ha[maskedByteLen]);
                    realByteNum += 1;
                }

            } else {
                System.arraycopy(Ha, 0, tempK, (i - 1) * 32, Ha.length);
                realByteNum += 32;
            }
        }

        byte[] k = new byte[realByteNum];
        System.arraycopy(tempK, 0, k, 0, realByteNum);
        return k;
    }

//    public static byte[] encrypt(byte[] publicKey, byte[] data) {
//        SM2PublicKey Q = new SM2PublicKey(new DERBitString(publicKey));
//        SM2Cipher cipherObj = encrypt(Q, data);
//        if (null == cipherObj) {
//            return null;
//        } else {
//            try {
//                return cipherObj.getEncoded();
//            } catch (IOException e) {
//                return null;
//            }
//        }
//    }
//
//    public static SM2Cipher encrypt(SM2PublicKey Q, byte[] m) {
//        if (Q == null) {
//            return null;
//        }
//
//        if (m == null || m.length == 0) {
//            return null;
//        }
//
//        SM2 sm2 = SM2.getInstance();
//        ECPoint P_B = sm2.getEccCurve().decodePoint(Q.getEncoded());
//
//        AsymmetricCipherKeyPair key = sm2.getEccKeyPairGenerator().generateKeyPair();
//        ECPrivateKeyParameters ecpriv = (ECPrivateKeyParameters) key.getPrivate();
//        ECPublicKeyParameters ecpub = (ECPublicKeyParameters) key.getPublic();
//        BigInteger k = ecpriv.getD();
//        ECPoint C1 = ecpub.getQ();
//
//        ECPoint S = P_B.multiply(sm2.getEccBcSpec().getH());
//        if (S.isInfinity()) {
//            return null;
//        }
//
//        ECPoint x2y2 = P_B.multiply(k);
//        byte[] x2 = x2y2.normalize().getXCoord().getEncoded();
//        byte[] y2 = x2y2.normalize().getYCoord().getEncoded();
//
//        byte[] z = new byte[x2.length + y2.length];
//        System.arraycopy(x2, 0, z, 0, x2.length);
//        System.arraycopy(y2, 0, z, x2.length, y2.length);
//
//        byte[] t = KDF(z, m.length * 8);
//
//        byte[] C2 = new byte[m.length];
//        System.arraycopy(m, 0, C2, 0, m.length);
//
//        for (int i = 0; i < C2.length; i++) {
//            C2[i] ^= t[i];
//        }
//
//
//        SM3Digest sm3 = new SM3Digest();
//        sm3.update(x2, 0, x2.length);
//        sm3.update(m, 0, m.length);
//        sm3.update(y2, 0, y2.length);
//
//        byte[] C3 = new byte[32];
//        sm3.doFinal(C3, 0);
//
//        ASN1Integer XCoordinate = new ASN1Integer(C1.normalize().getXCoord().toBigInteger());
//        ASN1Integer YCoordinate = new ASN1Integer(C1.normalize().getYCoord().toBigInteger());
//        ASN1OctetString HASH = new DEROctetString(C3);
//        ASN1OctetString CipherText = new DEROctetString(C2);
//        return new SM2Cipher(XCoordinate, YCoordinate, HASH, CipherText);
//    }

//    public static byte[] decrypt(byte[] privateKey, byte[] encryptedData) {
//        BigInteger key = new BigInteger(1, privateKey);
//
//        SM2Cipher c = SM2Cipher.getInstance(ASN1Sequence.getInstance(encryptedData));
//        SM2PrivateKey d = new SM2PrivateKey(new ASN1Integer(key));
//
//        return decrypt(d, c);
//    }

//    public static byte[] decrypt(SM2PrivateKey d, SM2Cipher c) {
//
//        BigInteger x = c.getXCoordinate().getPositiveValue();
//        BigInteger y = c.getYCoordinate().getPositiveValue();
//
//        SM2 sm2 = SM2.getInstance();
//        ECPoint C1 = sm2.getEccCurve().createPoint(x, y);
//        ECPoint S = C1.multiply(sm2.getEccBcSpec().getH());
//        if (S.isInfinity()) {
//            return null;
//        }
//
//        ECPoint x2y2 = C1.multiply(d.getKey().getPositiveValue());
//        byte[] x2 = x2y2.normalize().getXCoord().getEncoded();
//        byte[] y2 = x2y2.normalize().getYCoord().getEncoded();
//
//        byte[] z = new byte[x2.length + y2.length];
//        System.arraycopy(x2, 0, z, 0, x2.length);
//        System.arraycopy(y2, 0, z, x2.length, y2.length);
//
//        byte[] C2 = c.getCipherText().getOctets();
//        byte[] t = KDF(z, C2.length * 8);
//
//        byte[] m = new byte[C2.length];
//        System.arraycopy(C2, 0, m, 0, m.length);
//        for (int i = 0; i < m.length; i++) {
//            m[i] ^= t[i];
//        }
//
//
//        SM3Digest sm3 = new SM3Digest();
//        sm3.update(x2, 0, x2.length);
//        sm3.update(m, 0, m.length);
//        sm3.update(y2, 0, y2.length);
//
//        byte[] u = new byte[32];
//        sm3.doFinal(u, 0);
//
//        byte[] C3 = c.getHASH().getOctets();
//        if (Arrays.equals(u, C3)) {
//            return m;
//        } else {
//            return null;
//        }
//    }

    //	public static Sm2KeyPair generateKeyPair() {
    //
    //	}

    public static byte[] sign(byte[] privateKey, byte[] sourceData) {
        return sign(DEFAULT_SIGN_ID.getBytes(), privateKey, sourceData);
    }

    public static byte[] sign(byte[] userId, byte[] privateKey, byte[] sourceData) {
        SM2 sm2 = SM2.getInstance();

        BigInteger d_A = new BigInteger(1, privateKey);
        SM2PrivateKey d = new SM2PrivateKey(new ASN1Integer(d_A));
        ECPoint P_A = sm2.getEccPointG().multiply(d_A);

        SM2PublicKey Q = new SM2PublicKey(P_A);
        byte[] z = Z(userId, Q);
        byte[] H = H(z, sourceData);

        SM2Signature signData = sign(d, H);

        try {
            return signData.getEncoded();
        } catch (Exception e) {
            return null;
        }

    }

    private static SM2Signature sign(SM2PrivateKey d, byte[] H) {

        SM2 sm2 = SM2.getInstance();

        BigInteger k, x1;
        BigInteger r, s = null;
        ECPoint x1y1 = null;
        BigInteger e = new BigInteger(1, H);
        BigInteger d_A = d.getKey().getPositiveValue();
        boolean loop = false;
        do {
            AsymmetricCipherKeyPair key = sm2.getEccKeyPairGenerator().generateKeyPair();
            ECPrivateKeyParameters ecpriv = (ECPrivateKeyParameters) key.getPrivate();
            ECPublicKeyParameters ecpub = (ECPublicKeyParameters) key.getPublic();
            k = ecpriv.getD();
            x1y1 = ecpub.getQ();

            x1 = x1y1.normalize().getXCoord().toBigInteger();

            r = e.add(x1).mod(sm2.getEccN());
            if (BigInteger.ZERO.equals(r) || r.add(k).equals(sm2.getEccN())) {
                loop = true;
                continue;
            }

            BigInteger tmp1 = r.multiply(d_A);
            BigInteger tmp2 = k.subtract(tmp1);
            s = BigInteger.ONE.add(d_A).modInverse(sm2.getEccN()).multiply(tmp2).mod(sm2.getEccN());
            if (BigInteger.ZERO.equals(s)) {
                loop = true;
            }
        } while (loop);

        return new SM2Signature(new ASN1Integer(r), new ASN1Integer(s));
    }

//    public static boolean verifySign(byte[] publicKey, byte[] sourceData, byte[] signData) {
//        return verifySign(DEFAULT_SIGN_ID.getBytes(), publicKey, sourceData, signData);
//    }

    public static boolean verifySign(byte[] publicKey, byte[] sourceData, byte[] signData) {
        SM2Signature sign = SM2Signature.getInstance(signData);
        SM2PublicKey sm2PublicKey = new SM2PublicKey(new DERBitString(publicKey));
        byte[] H = generateSource(sourceData, sm2PublicKey);
        return verifySign(H, sign, sm2PublicKey);
    }

    public static boolean verifySign(Certificate cert, byte[] sourceData, byte[] signData) {
        if(signData.length == 64){
            signData = AsnUtil.asn1ToByte(IKICertUtils.getSM2Signature(signData));
        }
        SM2Signature sign = SM2Signature.getInstance(signData);
        SM2PublicKey sm2PublicKey = new SM2PublicKey(new DERBitString(CertTools.getSignPublicKey(cert)));
        byte[] H = generateSource(sourceData, sm2PublicKey);
        return verifySign(H, sign, sm2PublicKey);
    }

    /**
     * 生成签名原文
     *
     * @param data
     * @param publicKey
     * @return
     */
    public static byte[] generateSource(byte[] data, byte[] publicKey) {
        SM2PublicKey sm2PublicKey = new SM2PublicKey(new DERBitString(publicKey));
        return generateSource(data, sm2PublicKey);
    }

    /**
     * 生成签名原文
     *
     * @param data
     * @param publicKey
     * @return
     */
    public static byte[] generateSource(byte[] data, SM2PublicKey publicKey) {
        byte[] z = Z(DEFAULT_SIGN_ID.getBytes(), publicKey);
        return H(z, data);
    }

    private static boolean verifySign(byte[] H, SM2Signature sign, SM2PublicKey Q) {
        SM2 sm2 = SM2.getInstance();

        BigInteger r = sign.getR().getPositiveValue();
        BigInteger s = sign.getS().getPositiveValue();

        if (!SM2Utils.validSM2BigInteger(r) || !SM2Utils.validSM2BigInteger(s)) {
            return false;
        }

        BigInteger t = r.add(s).mod(sm2.getEccN());
        if (BigInteger.ZERO.equals(t)) {
            return false;
        }

        ECPoint P_A = sm2.getEccCurve().createPoint(Q.getXCoordinate().getPositiveValue(),
                Q.getYCoordinate().getPositiveValue());
        ECPoint x1y1 = sm2.getEccPointG().multiply(s).add(P_A.multiply(t));

        BigInteger e = new BigInteger(1, H);
        BigInteger R = e.add(x1y1.normalize().getXCoord().toBigInteger()).mod(sm2.getEccN());
        return R.equals(r);
    }
}
